import parser
import os
#import numpy as np
from .base import Base

class Caffe2ONNX(Base):
    def __init__(self, success_list, failed_list):
        print('Caffe2ONNX Init--------------------------')
        self.tool_path = os.path.join(parser.get_value('maca_converter_root'), parser.get_value('maca_converter_path'))
        self.model_path = os.path.join(parser.get_value('model_root'), parser.get_value('caffe_model_path'))
        self.success_list = success_list
        self.failed_list = failed_list
        self.input = os.path.join(parser.get_value('model_root'), parser.get_value('caffe_input_npy'))#parser.get_value('h5_input_npy')
        self.output = os.path.join(parser.get_value('model_root'), parser.get_value('caffe_output_npy'))#parser.get_value('h5_output_npy')


    def test(self):
        cwd = os.getcwd()
        output = os.path.join(cwd, 'tmp.onnx')

        cmd = 'python ' + 'model_convert.py' + ' --model_path ' + self.model_path + ' --model_type caffe ' + '--output ' + output
        
        print('test_caffe2onnx: ', cmd)

        '''
        data = np.load(self.input)
        input = np.ones_like(data)
        input[0][0] = data[0][2]
        input[0][1] = data[0][1]
        input[0][2] = data[0][0]

        path = self.input.split('.')[-2] + '_test.npy'
        np.save(path, input)
        '''

        self.test_model2onnx(self.tool_path, 'caffe2onnx', cmd, output, self.success_list, self.failed_list, self.input, self.output)